
#include "tds_protocol.h"
#include "log.h"
namespace secsmart_protocol
{
    DETECT_RESULT TdsProtocol::parse(secsmart_session::session_sptr session, const uint8_t *buf, size_t size, bool is_forward)
    {
        direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_UNKNOW;
        if (size < 8){ //小于协议长度
            return DETECT_RESULT::DETECT_UNKNOWN;
        }
        uint8_t type = (buf[0] & 0xff);
        LOG_DEBUG("type: {}", (int)type);
        if(!isTdsType(type))
        {
            if(type == TDS_SMP){//有Session Multiplex Protocol
                return parseSmp(session, buf, size, is_forward);
            }
            return DETECT_RESULT::DETECT_UNKNOWN;
        }
        uint16_t len = ((uint8_t)(buf[2]) << 8) + (uint8_t)(buf[3]);
        LOG_DEBUG("len = {}, data len= {}, direction = {}", len, size, (int)direction);
        if(len == (uint16_t)size && direction != secsmart_session::session_dirction_t::SESSION_DIRCTION_UNKNOW){
            
            if(!setIpPort(session, getDirction(), is_forward)){
                LOG_DEBUG("set port ip fail");
                return DETECT_RESULT::DETECT_FAIL;
            }
            addMatchCount();
            return DETECT_RESULT::DETECT_OK;
        }
        addNoMatchCount();
        return DETECT_RESULT::DETECT_UNKNOWN;  
    }

    bool TdsProtocol::isTdsType(uint8_t type)
    {
        bool result = false;
        switch(type){
            case TDS_SQL_BATCH:
            case TDS_LOGIN5:
            case TDS_RPC:
            case TDS_SIGNAL:
            case TDS_BULK_DATA:
            case TDS_TRANSACTION:
            case TDS_LOGIN7:
            case TDS_PRE_LOGIN:
            //请求的包类型
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                result = true;
                break;
            case TDS_QUERY5://sqlserver 没有这种类型，sybase有这种类型
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_REQUEST;
                databaseName = DBNAME_SYBASE;
                result = true;
                break;
            case TDS_RESPONSE:
            //返回的包类型
                direction = secsmart_session::session_dirction_t::SESSION_DIRCTION_RESPONSE;
                result = true;
                break;
            default:
                break;
        }
        return result;
    }

    DETECT_RESULT TdsProtocol::parseSmp(secsmart_session::session_sptr session, const uint8_t *buf, size_t size, bool is_forward)
    {
        if (size < 24){ //16字节的smp头长度，8字节tds头长度
            return DETECT_RESULT::DETECT_UNKNOWN;
        }
        uint8_t smid = (buf[0] & 0xff);
        if(TDS_SMP != smid){
            return DETECT_RESULT::DETECT_UNKNOWN;
        }
        uint8_t flag = (buf[1] & 0xff);
        switch(flag){
            case TDS_SMP_FLG_SYN:
            case TDS_SMP_FLG_ACK:
            case TDS_SMP_FLG_FIN:
            case TDS_SMP_FLG_DATA:
                break;
            default:
                return DETECT_RESULT::DETECT_UNKNOWN;
        }
        //接下去是2个字节的sessionid,跳过
        uint32_t len = (uint32_t)((buf[4] & 0xff) | ((buf[5] & 0xff) << 8) | ((buf[6] & 0xff) << 16) | ((buf[7] & 0xff) << 24));
        if(len != (uint32_t)size){
            return DETECT_RESULT::DETECT_UNKNOWN;
        }
        //4字节seqNum+4字节wndw

        uint8_t type = (buf[16] & 0xff); //tds包类型
        if(!isTdsType(type)){
            return DETECT_RESULT::DETECT_UNKNOWN;
        }
        len = (buf[18] & 0xff) << 8 + (buf[19] & 0xff);
        if(len == (uint16_t)(size - 16) && direction != secsmart_session::session_dirction_t::SESSION_DIRCTION_UNKNOW){
            if(!setIpPort(session, getDirction(), is_forward)){
                return DETECT_RESULT::DETECT_FAIL;
            }
            addMatchCount();
            return DETECT_RESULT::DETECT_OK;
        }
        addNoMatchCount();
        return DETECT_RESULT::DETECT_UNKNOWN;
    }

    std::string TdsProtocol::extendInfo(){
        return "";
    }
}